import copy

import torch

from algorithms.algorithm import Algorithm
from algorithms.large.mocgde import MOCGDE
from individual.population import Population
from individual.solution import Solution
from problems.problem import Problem


class MOEGS(Algorithm):
    parameters = [
        {'label': 'r', 'name': 'r', 'type': 'number', 'description': 'Number of test candidates', 'step': 1,
         'defaultValue': 500}
    ]

    def __init__(self, problem: Problem,
                 pop_size: int,
                 kwargs,
                 ):
        super().__init__(pop_size, problem, kwargs)
        self.population = problem.init_population(pop_size)
        self.archive = copy.deepcopy(self.population)
        self.thao = 0.1

    def final_output(self):
        return self.archive

    def each_iteration(self):
        for i in range(self.pop_size):
            parent = self.population[i]
            z = (torch.randn((self.problem.var_dim, self.r), dtype=torch.double) * self.thao).T
            offspring_dec = parent.dec + z
            offspring_dec = self.problem.repair_decision(offspring_dec)
            offspring_obj = self.problem.estimate_population(offspring_dec).obj
            offspring = Population(offspring_dec, offspring_obj)
            self.archive = MOCGDE.update_archive(self.archive + offspring, self.pop_size)
            if parent.con > 0:
                gk = torch.sum((offspring.con - parent.con), dim=1)
            elif offspring_obj.shape[1] == 2:
                w = torch.rand((self.r, 1))
                w = torch.cat([w, 1 - w], dim=1)
                gk = torch.sum(w * (offspring.obj - parent.obj), dim=1)
            elif offspring_obj.shape[1] == 3:
                w = torch.rand((self.r, 1))
                w = torch.cat([w / 3, (1 - w / 3) / 2, 1 - w / 3 - (1 - w / 3) / 2], dim=1)
                gk = torch.sum(w * (offspring.obj - parent.obj), dim=1)

            c = parent.dec - self.thao * torch.matmul(gk.unsqueeze(0), z) / torch.norm(torch.matmul(gk.unsqueeze(0), z))
            c = self.problem.repair_decision(c)
            offspring_obj = self.problem.estimate_population(c).obj
            offspring = Solution.solution_build(c, offspring_obj)
            if torch.sum(torch.max(offspring.con, torch.tensor(0.0))) < torch.sum(
                    torch.max(parent.con, torch.tensor(0.0))) or torch.all(offspring.obj < parent.obj):
                self.population[i] = offspring
                self.thao *= 1.8
            elif torch.sum(torch.max(parent.con, torch.tensor(0.0))) < torch.sum(
                    torch.max(offspring.con, torch.tensor(0.0))) or torch.all(parent.obj < offspring.obj):
                self.population[i] = self.archive[torch.randint(len(self.archive), (1,)).item()]
                self.thao /= 1.8
            else:
                if torch.rand(1) > 0.5:
                    self.population[i] = offspring
                else:
                    self.population[i] = self.archive[torch.randint(len(self.archive), (1,)).item()]
            self.archive = MOCGDE.update_archive(self.archive + offspring, self.pop_size)
            if self.thao > 1.2:
                self.thao /= 1.2
            elif self.thao < 1e-4:
                self.thao *= 1.2
